#############################################################################
# Windows Wallpaper Removal via MDM Personalization CSP
# Compatible with Windows PowerShell 5.1
# Removes enforced desktop/lock-screen wallpaper set via MDM Bridge
#############################################################################

# Step 0 - Force 64-bit PowerShell if running under 32-bit on x64 systems
if ($env:PROCESSOR_ARCHITEW6432 -eq "AMD64") {
    if ($myInvocation.Line) {
        &"$env:WINDIR\sysnative\WindowsPowerShell\v1.0\powershell.exe" -NonInteractive -NoProfile $myInvocation.Line
    } else {
        &"$env:WINDIR\sysnative\WindowsPowerShell\v1.0\powershell.exe" -NonInteractive -NoProfile -File "$($myInvocation.InvocationName)" $args
    }
    exit $LASTEXITCODE
}

#############################################################################
# Logging helper
#############################################################################
function Write-Log([string]$Message, [string]$Level = 'INFO') {
    $timestamp = Get-Date -Format "yyyy-MM-ddTHH:mm:ss"
    Write-Host "[$timestamp] [$Level] $Message"
}

#############################################################################
# Environment-based configuration
#############################################################################
function Get-BoolEnv {
    param([string]$Name, [bool]$Default)
    $v = [Environment]::GetEnvironmentVariable($Name)
    if ([string]::IsNullOrWhiteSpace($v)) { return $Default }
    switch -Regex ($v.Trim()) {
        '^(?i:true|1|yes|y)$' { return $true }
        '^(?i:false|0|no|n)$' { return $false }
        default { return $Default }
    }
}

$ClearDesktop    = Get-BoolEnv -Name 'CLEAR_DESKTOP'    -Default $true
$ClearLockScreen = Get-BoolEnv -Name 'CLEAR_LOCKSCREEN' -Default $true
$RegCleanup      = Get-BoolEnv -Name 'REG_CLEANUP'      -Default $true
$RestartExplorer = Get-BoolEnv -Name 'RESTART_EXPLORER' -Default $false

Write-Log "Wallpaper removal script initiated."
Write-Log "Options: CLEAR_DESKTOP=$ClearDesktop, CLEAR_LOCKSCREEN=$ClearLockScreen, REG_CLEANUP=$RegCleanup, RESTART_EXPLORER=$RestartExplorer"

#############################################################################
# Step 1 - Connect to MDM Bridge Provider
#############################################################################
$Namespace = 'root\cimv2\mdm\dmmap'
$ClassName = 'MDM_Personalization'

try {
    $Instance = Get-CimInstance -Namespace $Namespace -ClassName $ClassName -ErrorAction Stop
    Write-Log "Connected to MDM Bridge Provider successfully."
} catch {
    Write-Log "Failed to access MDM Bridge Provider. Run as SYSTEM. $_" 'ERROR'
    exit 4
}

#############################################################################
# Step 2 - Clear CSP values to remove wallpaper enforcement
#############################################################################
function Clear-CspProperty {
    param([string]$PropertyName)
    try {
        $script:Instance.$PropertyName = $null
        Set-CimInstance -CimInstance $script:Instance -ErrorAction Stop | Out-Null
        Write-Log "Cleared $PropertyName via null assignment (CSP Delete)."
        return $true
    } catch {
        Write-Log "Null clear failed for $PropertyName. Trying empty string..." 'WARN'
        try {
            $script:Instance.$PropertyName = ""
            Set-CimInstance -CimInstance $script:Instance -ErrorAction Stop | Out-Null
            Write-Log "Cleared $PropertyName via empty-string assignment."
            return $true
        } catch {
            Write-Log "Failed to clear $PropertyName. $($_.Exception.Message)" 'ERROR'
            return $false
        }
    }
}

$Changed = $false
if ($ClearDesktop) {
    if (Clear-CspProperty -PropertyName 'DesktopImageUrl') { $Changed = $true }
}
if ($ClearLockScreen) {
    if (Clear-CspProperty -PropertyName 'LockScreenImageUrl') { $Changed = $true }
}

if (-not $Changed) {
    Write-Log "No CSP values were changed; properties may already be empty." 'WARN'
}

#############################################################################
# Step 3 - Read back and verify CSP status
#############################################################################
function Get-StatusText([int]$code) {
    switch ($code) {
        1 {'Success'}
        2 {'InProgress'}
        3 {'Failed'}
        4 {'UnknownFileType'}
        5 {'UnsupportedUrlScheme'}
        6 {'MaxRetryFailed'}
        default {'Unknown'}
    }
}

try {
    $Instance = Get-CimInstance -Namespace $Namespace -ClassName $ClassName -ErrorAction Stop
    $dUrl = $Instance.DesktopImageUrl
    $lUrl = $Instance.LockScreenImageUrl
    $dStatus = $Instance.DesktopImageStatus
    $lStatus = $Instance.LockScreenImageStatus

    if ($ClearDesktop) {
        Write-Log ("DesktopImageUrl now='{0}' (Status={1} {2})" -f $dUrl, $dStatus, (Get-StatusText $dStatus))
    }
    if ($ClearLockScreen) {
        Write-Log ("LockScreenImageUrl now='{0}' (Status={1} {2})" -f $lUrl, $lStatus, (Get-StatusText $lStatus))
    }
} catch {
    Write-Log "Failed to read back MDM_Personalization: $($_.Exception.Message)" 'WARN'
}

#############################################################################
# Step 4 - Registry mirror cleanup (optional)
#############################################################################
if ($RegCleanup) {
    $RegPath = 'HKLM:\SOFTWARE\Microsoft\Windows\CurrentVersion\PersonalizationCSP'
    $Props = @()
    if ($ClearDesktop)    { $Props += 'DesktopImageUrl','DesktopImagePath','DesktopImageStatus' }
    if ($ClearLockScreen) { $Props += 'LockScreenImageUrl','LockScreenImagePath','LockScreenImageStatus' }

    if (Test-Path $RegPath) {
        foreach ($p in $Props | Select-Object -Unique) {
            try {
                if (Get-ItemProperty -Path $RegPath -Name $p -ErrorAction SilentlyContinue) {
                    Remove-ItemProperty -Path $RegPath -Name $p -ErrorAction Stop
                    Write-Log "Removed registry mirror value: $p"
                }
            } catch {
                Write-Log "Registry cleanup failed for $p. $($_.Exception.Message)" 'WARN'
            }
        }
    } else {
        Write-Log "Registry mirror path not found: $RegPath" 'INFO'
    }
}

#############################################################################
# Step 5 - Restart Explorer shell (optional)
#############################################################################
if ($RestartExplorer) {
    try {
        Write-Log "Restarting Explorer shell..."
        Get-Process explorer -ErrorAction SilentlyContinue | Stop-Process -Force -ErrorAction SilentlyContinue
        Start-Process explorer.exe
        Write-Log "Explorer restarted."
    } catch {
        Write-Log "Failed to restart Explorer. $($_.Exception.Message)" 'WARN'
    }
}

#############################################################################
# Step 6 - Completion
#############################################################################
Write-Log "Wallpaper policy removal completed successfully."
exit 0
